/*
 * A Demo to OpenCV Implementation of SURF
 * Further Information Refer to "SURF: Speed-Up Robust Feature"
 * Author: Liu Liu
 * liuliu.1987+opencv@gmail.com
 
 * g++ find_obj.cpp -o find `pkg-config --cflags --libs opencv`
 *
 */

#include <cv.h>
#include <highgui.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>

#include <iostream>
#include <vector>
#include <sys/types.h>
#include <dirent.h>
#include <errno.h>
#include <string>

using namespace std;

float pairs = 0;

IplImage *image = 0;
string fname;
float maxV, second, third = 0;
int indexMax, indexSec, indexTh = 0;
vector<string> files = vector<string>();

int getdir (string dir, vector<string> &files)
{
    DIR *dp;
    struct dirent *dirp;
    if((dp  = opendir(dir.c_str())) == NULL) {
        cout << "Error(" << errno << ") opening " << dir << endl;
        return errno;
    }

    while ((dirp = readdir(dp)) != NULL) {
        files.push_back(string(dirp->d_name));
    }
    closedir(dp);
    return 0;
}

double compareSURFDescriptors( const float* d1, const float* d2, double best, int length )
{
    double total_cost = 0;
    assert( length % 4 == 0 );
    for( int i = 0; i < length; i += 4 )
    {
        double t0 = d1[i] - d2[i];
        double t1 = d1[i+1] - d2[i+1];
        double t2 = d1[i+2] - d2[i+2];
        double t3 = d1[i+3] - d2[i+3];
        total_cost += t0*t0 + t1*t1 + t2*t2 + t3*t3;
        if( total_cost > best )
            break;
    }
    return total_cost;
}

int naiveNearestNeighbor( const float* vec, int laplacian,
                      const CvSeq* model_keypoints,
                      const CvSeq* model_descriptors ){
                      
    int length = (int)(model_descriptors->elem_size/sizeof(float));
    int i, neighbor = -1;
    double d, dist1 = 1e6, dist2 = 1e6;
    CvSeqReader reader, kreader;
    cvStartReadSeq( model_keypoints, &kreader, 0 );
    cvStartReadSeq( model_descriptors, &reader, 0 );

    for( i = 0; i < model_descriptors->total; i++ )
    {
        const CvSURFPoint* kp = (const CvSURFPoint*)kreader.ptr;
        const float* mvec = (const float*)reader.ptr;
    	CV_NEXT_SEQ_ELEM( kreader.seq->elem_size, kreader );
        CV_NEXT_SEQ_ELEM( reader.seq->elem_size, reader );
        if( laplacian != kp->laplacian )
            continue;
        d = compareSURFDescriptors( vec, mvec, dist2, length );
        if( d < dist1 )
        {
            dist2 = dist1;
            dist1 = d;
            neighbor = i;
        }
        else if ( d < dist2 )
            dist2 = d;
    }
    if ( dist1 < 0.6*dist2 )
        return neighbor;
    return -1;
}

void findPairs( const CvSeq* objectKeypoints, const CvSeq* objectDescriptors,
           const CvSeq* imageKeypoints, const CvSeq* imageDescriptors, vector<int>& ptpairs )
{
    int i;
    CvSeqReader reader, kreader;
    cvStartReadSeq( objectKeypoints, &kreader );
    cvStartReadSeq( objectDescriptors, &reader );
    ptpairs.clear();

    for( i = 0; i < objectDescriptors->total; i++ )
    {
        const CvSURFPoint* kp = (const CvSURFPoint*)kreader.ptr;
        const float* descriptor = (const float*)reader.ptr;
        CV_NEXT_SEQ_ELEM( kreader.seq->elem_size, kreader );
        CV_NEXT_SEQ_ELEM( reader.seq->elem_size, reader );
        int nearest_neighbor = naiveNearestNeighbor( descriptor, kp->laplacian, imageKeypoints, imageDescriptors );
        if( nearest_neighbor >= 0 )
        {
            ptpairs.push_back(i);
            ptpairs.push_back(nearest_neighbor);
        }
    }
}

/* a rough implementation for object location */
int locatePlanarObject( const CvSeq* objectKeypoints, const CvSeq* objectDescriptors,
                    const CvSeq* imageKeypoints, const CvSeq* imageDescriptors,
                    const CvPoint src_corners[4], CvPoint dst_corners[4] ){
    double h[9];
    CvMat _h = cvMat(3, 3, CV_64F, h);
    vector<int> ptpairs;
    vector<CvPoint2D32f> pt1, pt2;
    CvMat _pt1, _pt2;
    int i, n;

    findPairs( objectKeypoints, objectDescriptors, imageKeypoints, imageDescriptors, ptpairs );

    n = ptpairs.size()/2;
    pairs = (float)n;
    if( n < 4 )
        return 0;

    pt1.resize(n);
    pt2.resize(n);
    for( i = 0; i < n; i++ )
    {
        pt1[i] = ((CvSURFPoint*)cvGetSeqElem(objectKeypoints,ptpairs[i*2]))->pt;
        pt2[i] = ((CvSURFPoint*)cvGetSeqElem(imageKeypoints,ptpairs[i*2+1]))->pt;
    }

    _pt1 = cvMat(1, n, CV_32FC2, &pt1[0] );
    _pt2 = cvMat(1, n, CV_32FC2, &pt2[0] );
    if( !cvFindHomography( &_pt1, &_pt2, &_h, CV_RANSAC, 5 ))
        return 0;

    for( i = 0; i < 4; i++ )
    {
        double x = src_corners[i].x, y = src_corners[i].y;
        double Z = 1./(h[6]*x + h[7]*y + h[8]);
        double X = (h[0]*x + h[1]*y + h[2])*Z;
        double Y = (h[3]*x + h[4]*y + h[5])*Z;
        dst_corners[i] = cvPoint(cvRound(X), cvRound(Y));
    }

    return 1;
}

float surf(IplImage *image, IplImage *object, int count){
	CvMemStorage* storage = cvCreateMemStorage(0);
	
	CvSeq *objectKeypoints = 0, *objectDescriptors = 0;
	CvSeq *imageKeypoints = 0, *imageDescriptors = 0;

	CvSURFParams params = cvSURFParams(500, 1);

	double tt = (double)cvGetTickCount();
	cvExtractSURF( object, 0, &objectKeypoints, &objectDescriptors, storage, params );
	//printf("Object Descriptors: %d\n", objectDescriptors->total);
	
	cvExtractSURF( image, 0, &imageKeypoints, &imageDescriptors, storage, params );
	//printf("Image Descriptors: %d\n", imageDescriptors->total);
	
	tt = (double)cvGetTickCount() - tt;
	//printf( "Extraction time = %gms\n", tt/(cvGetTickFrequency()*1000.));
	CvPoint src_corners[4] = {{0,0}, {object->width,0}, {object->width, object->height}, {0, object->height}};
	CvPoint dst_corners[4];
	IplImage* correspond = cvCreateImage( cvSize(image->width, object->height+image->height), 8, 1 );
	cvSetImageROI( correspond, cvRect( 0, 0, object->width, object->height ) );
	cvCopy( object, correspond );
	cvSetImageROI( correspond, cvRect( 0, object->height, correspond->width, correspond->height ) );
	cvCopy( image, correspond );
	cvResetImageROI( correspond );

	locatePlanarObject( objectKeypoints, objectDescriptors, imageKeypoints,imageDescriptors, src_corners, dst_corners );

	vector<int> ptpairs;
	findPairs( objectKeypoints, objectDescriptors, imageKeypoints, imageDescriptors, ptpairs );

	if(pairs <= 2)
		pairs = 0;
	
	return pairs/objectDescriptors->total;
}

int main(int argc, char** argv){
    string dir = string("./image_recognition/objects/objects_black/");
    getdir(dir,files);

	const char* scene_filename = argc == 2 ? argv[1] : "./image_recognition/scenes/test_126_cam0.jpg";

	/////////////////////////////////////////////////
	/////										                    /////
	/////			FIND SHAPE				                /////
	/////////////////////////////////////////////////
	
	for (int i = 0; i < (int)files.size(); i++) {
	   if(files[i].find(".jpg") == string::npos)
	   		continue;

		fname = scene_filename;

		IplImage* object = cvLoadImage( (dir + files[i]).c_str(), CV_LOAD_IMAGE_GRAYSCALE );
		IplImage* image = cvLoadImage( scene_filename, CV_LOAD_IMAGE_GRAYSCALE );
	
		if(!image ){
			cout<<"Image!\n";
			exit(-1);
		}
		if(!object){
			cout<<"Object!\n";
			exit(-1);
		}
	
		cvSmooth(object, object, CV_GAUSSIAN, 13);
		cvSmooth(image, image, CV_GAUSSIAN, 1);
	  
		float wtf = surf(image, object, i);
		printf("%s   \t%f\n", files[i].c_str(), wtf); 
	
		if(wtf > maxV){
			maxV = wtf;
			indexMax = i;
		}	
  	}  
  	
  	printf("\nShape is %s with value %f.\n", files[indexMax].c_str(), maxV);
  	maxV = 0;

	/////////////////////////////////////////////////
	/////										/////
	/////			FIND COLOR	              / ////
	/////////////////////////////////////////////////
	if(files[indexMax].find("apple") != string::npos)
		dir = string("./image_recognition/objects/dryer/");
	if(files[indexMax].find("banana") != string::npos)
		dir = string("./image_recognition/objects/banana/");
	if(files[indexMax].find("book") != string::npos)
		dir = string("./image_recognition/objects/book/");
	if(files[indexMax].find("camera") != string::npos)
		dir = string("./image_recognition/objects/camera/");
	if(files[indexMax].find("computer") != string::npos)
		dir = string("./image_recognition/objects/computer/");
	if(files[indexMax].find("dryer") != string::npos)
		dir = string("./image_recognition/objects/dryer/");
	if(files[indexMax].find("glass") != string::npos)
		dir = string("./image_recognition/objects/glass/");
	if(files[indexMax].find("glasses") != string::npos)
		dir = string("./image_recognition/objects/glasses/");
	if(files[indexMax].find("hammer") != string::npos)
		dir = string("./image_recognition/objects/hammer/");		
	if(files[indexMax].find("mug") != string::npos)
		dir = string("./image_recognition/objects/mug/");		
	if(files[indexMax].find("scissors") != string::npos)
		dir = string("./image_recognition/objects/scissors/");
	if(files[indexMax].find("teddy") != string::npos)
		dir = string("./image_recognition/objects/teddy/");
	
	vector<string> files = vector<string>();
	getdir(dir,files);
	indexMax = 0;
	for (int i = 0; i < (int)files.size(); i++) {
	   if(files[i].find(".jpg") == string::npos)
	   		continue;

		IplImage* object = cvLoadImage( (dir + files[i]).c_str(), CV_LOAD_IMAGE_UNCHANGED );
		IplImage* image = cvLoadImage( scene_filename, CV_LOAD_IMAGE_UNCHANGED );
	
		if(!image ){
			cout<<"Image!\n";
			exit(-1);
		}
		if(!object){
			cout<<"Object!\n";
			exit(-1);
		}
	
		cvSmooth(object, object, CV_GAUSSIAN, 13);
		cvSmooth(image, image, CV_GAUSSIAN, 1);

		CvSize s = cvSize(object->width, object->height);
		int d = object->depth;
		IplImage* Ro = cvCreateImage(s, d, 1);
		IplImage* Go = cvCreateImage(s, d, 1);
		IplImage* Bo = cvCreateImage(s, d, 1);
		cvSplit(object, Bo, Go, Ro, 0);
	
		CvSize s2 = cvSize(image->width, image->height);
		d = image->depth;
		IplImage* Ri = cvCreateImage(s2, d, 1);
		IplImage* Gi = cvCreateImage(s2, d, 1);
		IplImage* Bi = cvCreateImage(s2, d, 1);
		cvSplit(image, Bi, Gi, Ri, 0);
	  
		float wtf = sqrt(pow(surf(Ri, Ro, i), 2) + pow(surf(Gi, Go, i), 2) + pow(surf(Bi, Bo, i), 2));
		printf("%s   \t%f\n", files[i].c_str(), wtf); 
	
		if(wtf > maxV){
			third = second;
			second = maxV;
			indexTh = indexSec;
			indexSec = indexMax;
	
			maxV = wtf;
			indexMax = i;
		}
		else if(wtf > second){
			third = second;
			indexTh = indexSec;
	
			second = wtf;
			indexSec = i;
		}
		else if(wtf > third){
			third = wtf;
			indexTh = i;
		}
  }  		
    
  printf("\n1: %f %s\n2: %f %s\n3: %f %s\n", maxV, files[indexMax].c_str(), second, files[indexSec].c_str(), third, files[indexTh].c_str());
  
  if((maxV - second)/maxV > 0.08)
      printf("Tag is %s!\n", files[indexMax].c_str());
  else{
    int pos = (int)files[indexMax].find("_");
    string tag = files[indexMax].substr(0, pos);
    
    
    printf("Tag is %s!\n", tag.c_str());
  }
  return 0;
}

